#pragma once

#include <utility>
#include <vector>

#include "planner/group_expression.h"
#include "planner/memo.h"
#include "planner/optimizer_task_pool.h"
#include "planner/rule.h"

namespace optimizer {

class OptimizerContext {
 public:
  explicit OptimizerContext() = default;

  ~OptimizerContext() {
    for (auto *ctx : track_list_) {
      delete ctx;
    }
  }

  Memo &GetMemo() { return memo_; }

  RuleSet &GetRuleSet() { return rule_set_; }

  OptimizerTaskStack &GetTaskPool() { return task_pool_; }

  void AddOptimizationContext(OptimizationContext *ctx) { track_list_.push_back(ctx); }

  void PushTask(OptimizerTask *task) { task_pool_.Push(task); }

  GroupExpression *MakeGroupExpression(std::shared_ptr<Operator> node);

  bool RecordOptimizerNodeIntoGroup(std::shared_ptr<Operator> node, GroupExpression **gexpr) {
    return RecordOptimizerNodeIntoGroup(std::move(node), gexpr, UNDEFINED_GROUP);
  }

  bool RecordOptimizerNodeIntoGroup(std::shared_ptr<Operator> node, GroupExpression **gexpr, group_id_t target_group) {
    auto *new_gexpr = MakeGroupExpression(std::move(node));
    auto *ptr = memo_.InsertExpression(new_gexpr, target_group, false);
    CASCADES_ASSERT(ptr, "Root of expr should not fail insertion");

    (*gexpr) = ptr;
    return (ptr == new_gexpr);
  }

  void ReplaceRewriteExpression(std::shared_ptr<Operator> node, group_id_t target_group) {
    memo_.EraseExpression(target_group);
    UNUSED_ATTRIBUTE auto *ret = memo_.InsertExpression(MakeGroupExpression(std::move(node)), target_group, false);
    CASCADES_ASSERT(ret, "Root expr should always be inserted");
  }

 private:
  Memo memo_;
  RuleSet rule_set_;
  OptimizerTaskStack task_pool_;
  std::vector<OptimizationContext *> track_list_;
};

}  // namespace optimizer
